from dataclasses import dataclass,asdict, field
from dataclasses_json import dataclass_json
from typing import List
import json 
import os
import numpy as np
from bs37ToARKit import ConvertBs37ToARKit51
from pipeline.utils.bs61 import *
from preprocess import bs_preprocess

@dataclass_json
@dataclass
class Vector3():
    x:float=0
    y:float=0
    z:float=0
        
@dataclass_json
@dataclass
class Vector4():
    x:float=0
    y:float=0
    z:float=0
    w:float=0
@dataclass_json
@dataclass
class Frame():
    location:list = field(default_factory=lambda : [])
    weight:list = field(default_factory=lambda : [])
    time:float=0
    position:list = field(default_factory=lambda : [])
    rotation:list = field(default_factory=lambda : [])

@dataclass_json
@dataclass
class RecordDataJson():   
    frames:list = field(default_factory=lambda : [])

def fromBEATJson(srcPath,dstPath):
    blendshape = []
    with open(srcPath,'r',encoding='utf8')as fp:
        file_contents = fp.read()                
        parsed_json = json.loads(file_contents)
        for j in range(len(parsed_json['frames'])):
            if(j % 2 == 0):
                item = parsed_json['frames'][j]
                weights = []
                for num in item['weights']:
                    weights.append(num)

                blendshape.append(weights)
    blendshape = np.array(blendshape)
    
    blendshape = bs_preprocess(blendshape)
    
    fn = os.path.basename(srcPath)
    fn = fn.replace('.json','.txt')
    np.savetxt(os.path.join(dstPath, fn), blendshape, fmt='%.8f')

def from3DETFNpy(srcPath,dstPath):

    blendshape = np.load(srcPath)
    blendshape = blendshape[:,:-1]

    fn = os.path.basename(srcPath)
    fn = fn.replace('.npy','.txt')
    np.savetxt(os.path.join(dstPath, fn), blendshape, fmt='%.8f')


def tensorToUnityBSJson(input,bs_scale=1.0,bs_delay=0.0):
    rdj = RecordDataJson()
    time = 0
    step = 1.0 / 30
       
    
    
    for i in range(input.shape[0]):
        frame = Frame()
        frame.time = time + bs_delay
        time = time + step
        frame.location = list(range(52))

        weights = []
        sweights = input[i]
        for sw in sweights:
            weights.append((float)(sw) * 100 * bs_scale)
        
        if len(weights) == 37:
            weights = ConvertBs37ToARKit51(weights)

        weights.append(0.0)#last tounge bs
        frame.weight = weights
        rdj.frames.append(frame)
    return rdj

def txtToUnityBSJson(srcTxtPath,bs_scale=1.0,bs_delay=0.0):
    

    rdj = RecordDataJson()
    time = 0
    step = 1.0 / 30

    with open(srcTxtPath) as f:
        lines = f.readlines()
        for line in lines:
            frame = Frame()
            frame.time = time + bs_delay
            time = time + step
            frame.location = list(range(52))

            weights = []
            sweights = line.split(' ')
            for sw in sweights:
                weights.append((float)(sw) * 100 * bs_scale)
            
            if len(weights) == 37:
                weights = ConvertBs37ToARKit51(weights)

            weights.append(0.0)#last tounge bs
            frame.weight = weights
            rdj.frames.append(frame)
    
    
    jstr = json.dumps(asdict(rdj))
    print(jstr)
    
    jfilePath = srcTxtPath + '.bs_weight.json'
    f2 = open(jfilePath, 'w')
    f2.write(jstr)
    f2.close()
    
def neurosyncToUnityJson(srcPath,bs_scale = 1.0):
    for f in os.listdir(srcPath):
        if f.endswith('npy'):
            input = np.load(os.path.join(srcPath,f))
                
            rdj = RecordDataJson()
            time = 0
            step = 1.0 / 60
            
            
            
            for i in range(input.shape[0]):
                frame = Frame()
                frame.time = time
                time = time + step
                frame.location = list(range(52))

                weights = []
                sweights = input[i]
                
                for j in range(51):
                    weight = input[i,map_52_to_61(j)]
                    weights.append((float)(weight) * 100 * bs_scale)
                
                weights.append(0.0)#last tounge bs
                frame.weight = weights
                rdj.frames.append(frame)
            
            
            jstr = json.dumps(asdict(rdj))
            print(jstr)
            
            jfilePath = os.path.join(srcPath,f + '.bs_weight.json') 
            f2 = open(jfilePath, 'w')
            f2.write(jstr)
            f2.close()
            
            
    
if __name__ == '__main__':

 
    # srcBeatJsonPath = 'data/test/1_wayne_0_1_1.json'
    # srcBeatJsonPath = 'data/test/12_zhao_2_57_64.json'
    # # srcBeatJsonPath = 'D:/DevelopProj/Yuji/Audio2Face/Audio2Face/data/test/10_kieks_0_73_74.json'
    # # srcBeatJsonPath = 'D:/DevelopProj/Yuji/Audio2Face/Audio2Face/data/test/11_nidal_0_100_100.json'
    # dstPath = 'D:/DevelopProj/Yuji/Audio2Face/YAudio2BS/data/pipeline/synthesis/'
    # # fromBEATJson(srcBeatJsonPath,dstPath)

    # src3DETFNpyPath = 'D:/Beat/3D-ETF/3D-ETF/all/WRA_JoniErnst0_000.npy'
    # from3DETFNpy(src3DETFNpyPath,dstPath)
   
    # srcTxtPath = os.path.join(dstPath,'1_wayne_0_1_1.txt')
    # srcTxtPath = os.path.join(dstPath,'12_zhao_2_57_64.txt')
    # srcTxtPath = os.path.join(dstPath,'WRA_JoniErnst0_000.txt')
    # # srcTxtPath = 'D:/DevelopProj/Yuji/Audio2Face/Audio2Face/synthesis/10_kieks_0_73_74.txt'
    # # srcTxtPath = 'D:/DevelopProj/Yuji/Audio2Face/Audio2Face/synthesis/pad-smooth-test1_wayne_0_98_98-lstmae-2distconcat_kl0.txt'
    # # srcTxtPath = 'D:/DevelopProj/Yuji/Audio2Face/Audio2Face/synthesis/pad-smooth-test18_daiki_0_8_8-lstmae-2distconcat_kl0.txt'
    # # srcTxtPath = 'D:/DevelopProj/Yuji/Audio2Face/Audio2Face/synthesis/pad-smooth-testwoman-lstmae-2distconcat_kl0.txt'
   
    # # srcTxtPath = dstPath + '1_wayne_0_100_100.txt'
    # txtToUnityBSJson(srcTxtPath)

    neurosyncPath = "D:/DevelopProj/Yuji/Audio2Face/YAudio2BS/data/neurosync"
    neurosyncToUnityJson(neurosyncPath)
